from sklearn.datasets import fetch_openml
from sklearn.preprocessing import StandardScaler
from torch import nn
from torch import optim
import torch
from sklearn.model_selection import train_test_split


class LinearRegressionModel(nn.Module):
    def __init__(self, input_dim: int):
        super(LinearRegressionModel, self).__init__()
        self.model = nn.Linear(input_dim, 1)
        self.loss = nn.MSELoss()
        self.scaler = StandardScaler()
        self.optimizer = None
        self.X_train = None
        self.X_test = None
        self.Y_train = None
        self.Y_test = None

    def load_data(self, origin_data: tuple, train_size: float = 0.2):
        data, target = origin_data
        # print(data)
        # print(target)
        x_train, x_test, y_train, y_test = train_test_split(data, target, test_size=train_size, random_state=42)

        print(x_train)

        x_train = self.scaler.fit_transform(x_train)
        x_test = self.scaler.transform(x_test)

        self.X_train = torch.tensor(x_train, dtype=torch.float32)
        self.X_test = torch.tensor(x_test, dtype=torch.float32)
        self.Y_train = torch.tensor(y_train, dtype=torch.float32).reshape(-1, 1)
        self.Y_test = torch.tensor(y_test, dtype=torch.float32).reshape(-1, 1)

    def train_model(self, epochs: int, lr: float = 0.01):
        self.optimizer = optim.SGD(self.model.parameters(), lr=lr)
        for epoch in range(epochs):
            outputs = self.model(self.X_train)
            loss = self.loss(outputs, self.Y_train)
            # 反向传播和优化
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            if (epoch + 1) % 10 == 0:
                print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}')

        # 6. 评估模型
        self.model.eval()
        with torch.no_grad():
            predicted = self.model(self.X_test)
            print(predicted)
            mse = self.loss(predicted, self.Y_test)
            print(f'Mean Squared Error on test set: {mse.item():.4f}')


if __name__ == '__main__':
    boston = fetch_openml(name="boston", version=1, as_frame=False)
    data, target = boston.data, boston.target
    model = LinearRegressionModel(data.shape[1])
    model.load_data((data, target))
    model.train_model(200)

    print(model.model(torch.randn((1, 13))))

